import json
import os
import argparse
from time import sleep
import tempfile

from watchfiles import awatch

from config_parse_util import SmokeCaseConfigs
from check_npu_status import check_multi_npu_not_used
from get_versions import get_version_from_pyproject, process_version
from inventory_template_tool import Inventory

script_path = os.path.dirname(os.path.abspath(__file__))
version_tmp_root_dir = "/data/CI/scripts/code_versions/"
"""
    1、读取配置文件，获取关键信息
    2、进入镜像，检查版本
    3、读取用例列表
     - 先从inventory或main_node中找到执行节点，检查是否npu空闲，如果没用空闲，等待3分钟再检查，直到连续两次确认空闲状态(避免是执行的中间状态)
     - 确认空闲后启动执行用例
"""
DEFAULT_REPOSITORY = ""

class Config:
    def __init__(self, config_path: str, config_root_dir: str, ver_conf_template: str, version_mode: bool):
        self.config_path = config_path
        self.config_root_dir = config_root_dir
        self.ver_conf_template = ver_conf_template
        self.version_mode = version_mode

    def generate_config_path(self, version: str):
        config_path = self.config_path
        if self.version_mode:
            if not version:
                raise ValueError("version mode, version is required")
            config_path = os.path.join(self.config_root_dir, version, self.ver_conf_template)
        if not os.path.exists(config_path):
            raise ValueError("config path: {config_path} does not exist".format(config_path=config_path))
        return config_path


class NodeManagement:
    def __init__(self, config_path: str, pool_id: str):
        self.config_path = config_path
        self.pool_id = pool_id
        self.pool_config = {}
        self.host_list = []
        self.idle_host_list = []
        self.busy_host_list = []
        self.init_machines()

    def init_machines(self):
        with open(self.config_path, "r") as f:
            config_list = json.load(f)
            for config in config_list:
                if config['id'] == self.pool_id:
                    self.pool_config = config
                    break
            if self.pool_config == {}:
                self.pool_config = config_list[0]
        self.host_list = self.pool_config["host_list"]
        check_npu_not_use, idle_hosts, busy_hosts = check_multi_npu_not_used(host_list=self.host_list,
                                                                             username=self.pool_config["login_user"],
                                                                             password=self.pool_config["login_password"],
                                                                             pkey=self.pool_config["ssh_private_key"])
        self.idle_host_list = idle_hosts
        self.busy_host_list = busy_hosts
        print(f"空闲节点: {self.idle_host_list}, 忙碌节点: {self.busy_host_list}")


class SmokeFramework:
    def __init__(self, _config: Config, image_name: str, output_path: str, parallel: int, dynamic_mode: bool, _node_manager: NodeManagement):
        self.image_name = image_name
        self.output_path = output_path
        self.parallel = parallel
        self.ssh_private_key = ""
        self.script_path = ""
        self.code_version = ""
        self.case_list = []
        self.get_code_version_from_image()  # 获取代码版本
        self.config_path = _config.generate_config_path(self.code_version)
        self.parse_config()
        os.makedirs(output_path, exist_ok=True)
        self.config_file_name = os.path.basename(self.config_path).split(".")[0]
        self.run_result_file = os.path.join(self.output_path, f"{self.config_file_name}_result.txt")
        self.dynamic_mode = dynamic_mode
        self.node_manager = _node_manager

    def parse_config(self):
        with open(self.config_path, 'r') as f:
            config_content = json.load(f)
            self.ssh_private_key = config_content['ssh_private_key']
            self.script_path = config_content['script_path']
            self.case_list = config_content['cases']

    def get_code_version_from_image(self):
        # 从镜像中获取代码版本
        os.makedirs(version_tmp_root_dir, exist_ok=True)
        with tempfile.TemporaryDirectory(dir=version_tmp_root_dir) as temp_dir:
            print("临时目录路径:", temp_dir)
            # 在此目录中操作文件...
            try:
                print(f"拉取镜像：{self.image_name}")
                os.system(f"docker pull {self.image_name}")
                print("启动容器，复制文件...")
                cp_cmd = f"cp /workspace/omniinfer/pyproject.toml {temp_dir}"

                cmd = ("docker run --rm --name \"take_omni_files\""
                       " -u root --privileged=true --entrypoint=bash --net=host"
                       f" -w {temp_dir} -v {version_tmp_root_dir}:{version_tmp_root_dir}"
                       f" -v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime:ro"
                       f" {self.image_name}"
                       f" -c \"source ~/.bashrc && {cp_cmd}\"")

                print(cmd)
                os.system(cmd)
                print("容器释放.")
            except Exception as e:
                print(f"docker get version fail, Error: {e}")
            # 读取版本
            temp_version_file = os.path.join(temp_dir, "pyproject.toml")
            original_version = get_version_from_pyproject(temp_version_file)

            if original_version:
                processed_version = process_version(original_version)
                self.code_version = processed_version
            print(f"当前代码版本：{processed_version}, 执行用例版本：{self.code_version}")

    def continue_check_npu(self, all_nodes: list):
        unused_npu_sum_times = 0
        while True:
            check_npu_not_use, idle_hosts, busy_hosts = check_multi_npu_not_used(host_list=all_nodes,
                                                                                 username="root",
                                                                                 password="",
                                                                                 pkey=self.ssh_private_key)
            if not check_npu_not_use:
                # npu占用，等待3分钟，继续检查
                print("目标节点npu占用，等待3min重复检查")
                sleep(180)
                continue
            else:
                # npu未占用，等待3分钟，多检查一次，当检查两次都没有使用时认为是空闲状态
                unused_npu_sum_times += 1
                # todo 检查是否有docker镜像
                print("目标节点npu暂无占用，开始执行用例")
                break

    def retry_fail_cases(self):
        # 只执行失败的用例
        tmp_case_list = []
        with open(self.run_result_file, "r") as f:
            while True:
                result_line = f.readline()
                if not result_line:
                    break
                result = result_line.replace(" ", "").split("|")
                if result[2] == "FAILED":
                    tmp_case_list.append(result[0])
        new_config_list = []
        for case_config in self.case_list:
            if case_config['case_num'] in tmp_case_list:
                new_config_list.append(case_config)
        return new_config_list

    def run_smoke(self, case_config: dict):
        open_intercept = case_config["open_intercept"]
        case_config['case_detail'] = "\"" + case_config['case_detail'] + "\""
        cmd = (f"cd {self.script_path} && bash rdv_smoke_runner.sh {case_config['case_num']} "
               f"{case_config['case_id']} {case_config['case_detail']} {case_config['case_run_script']}"
               f" {case_config['case_remove_scripts']} {case_config['run_performance']}"
               f" {case_config['run_precision']} {case_config['code_version']}")
        ret_code = os.system(cmd)
        if not open_intercept:
            print(f"{case_config['case_num']}: {case_config['case_detail']}仅执行验证，不进行拦截")
            return
        with open(self.run_result_file, "a") as f:
            if ret_code == 0:
                f.write(f"{case_config['case_num']}  |  {case_config['case_detail']}  |  SUCCESS")
                f.write("\n")
            else:
                f.write(f"{case_config['case_num']}  |  {case_config['case_detail']}  |  FAILED")
                f.write("\n")

    def get_parsed_nodes(self, parser: SmokeCaseConfigs):
        all_nodes = []
        parser.insert_hosts()
        node_configs = parser.configs
        for node in node_configs:
            if node['name'] == 'C0':
                continue
            all_nodes.append(node['host'])
        # 检查节点npu
        # all_nodes_str = ",".join(all_nodes)
        self.continue_check_npu(all_nodes)

    def run_cases_static_nodes(self, case_config: dict):
        """
            静态节点模式
            :param case_config: 用例配置
        """
        if "inventory" in case_config.keys():
            # 解析inventory文件
            parser = SmokeCaseConfigs(case_config["inventory"], "", "")
            self.get_parsed_nodes(parser)
            # 执行用例
            self.run_smoke(case_config)
        elif "hybrid_node" in case_config.keys():
            main_node = case_config["hybrid_node"]
            # 检查节点npu
            self.continue_check_npu(list(main_node))
            # 执行用例
            self.run_smoke(case_config)
        elif "multi_nodes" in case_config.keys():
            multi_nodes = case_config["multi_nodes"]
            # 检查节点npu
            self.continue_check_npu(multi_nodes)
            # 执行用例
            self.run_smoke(case_config)
        elif "cli" in case_config.keys():
            # 解析inventory文件
            parser = SmokeCaseConfigs("", "", case_config["cli"])
            self.get_parsed_nodes(parser)
            # 执行用例
            self.run_smoke(case_config)
        else:
            return

    def run_cases_dynamic_nodes(self, case_config: dict):
        """
            动态节点模式
            :param case_config: 用例配置
        """
        if "inventory" in case_config.keys():
            # 解析inventory文件
            parser = SmokeCaseConfigs(case_config["inventory"], "", "")
            self.get_parsed_nodes(parser)
            # 执行用例
            self.run_smoke(case_config)
        elif "hybrid_node" in case_config.keys():
            main_node = case_config["hybrid_node"]
            # 检查节点npu
            self.continue_check_npu(list(main_node))
            # 执行用例
            self.run_smoke(case_config)
        elif "multi_nodes" in case_config.keys():
            multi_nodes = case_config["multi_nodes"]
            # 检查节点npu
            self.continue_check_npu(multi_nodes)
            # 执行用例
            self.run_smoke(case_config)
        elif "cli" in case_config.keys():
            # 解析inventory文件
            parser = SmokeCaseConfigs("", "", case_config["cli"])
            self.get_parsed_nodes(parser)
            # 执行用例
            self.run_smoke(case_config)
        else:
            return

    def run_cases_parallel(self, case_config: dict):
        # 检查是否执行
        if not case_config["open_execute"]:  # 未打开, 跳过执行
            print(f"跳过执行{case_config['case_num']}: {case_config['case_detail']}")
            return

        # 当用例层级存在
        if case_config["code_version"]:
            case_code_version = case_config["code_version"]
            case_code_version_list = case_code_version.split(",")
            if self.code_version not in case_code_version_list or self.code_version not in case_code_version:
                print(f"版本不一致，跳过执行{case_config['case_num']}: {case_config['case_detail']}")
                return

        if self.dynamic_mode and self.node_manager is not None:   # 动态节点模式
            self.run_cases_dynamic_nodes(case_config)
        else:   # 静态节点模式
            self.run_cases_static_nodes(case_config)

    def run(self):
        if os.path.exists(self.run_result_file):
            self.case_list = self.retry_fail_cases()
            print(f"失败重试！需要重新执行如下用例：{self.case_list}")
            pass

        for case_config in self.case_list:
            self.run_cases_parallel(case_config)

        # 用例执行结束，打印结果，判断是否通过
        final_code = 0
        print("---------------- 执行结果 ----------------")
        with open(self.run_result_file, "r") as f:
            while True:
                line = f.readline()
                if not line:
                    break
                print(line)
                if "FAILED" in line:
                    final_code = 1
        print("----------------------------------------")
        return final_code

if __name__ == '__main__':
    """
    直接配置模式
    python smoke_framework_run.py --config_path /data/CI/scripts/omni_infer_smoke_framework/smoke_configs/smoke_case_config_tmp.json --image_name swr.cn-east-4.myhuaweicloud.com/omni-ci/rdv-arm-euler-py311-910c:20251021_792 --output_path /data/zkm/tmp
    版本配置模式
    python smoke_framework_run.py --config_root_dir /data/CI/scripts/omni_infer_smoke_framework/smoke_configs --ver_conf_template  smoke_case_config.json --image_name swr.cn-east-4.myhuaweicloud.com/omni-ci/rdv-arm-euler-py311-910c:20251021_792 --output_path /data/zkm/tmp --version_mode
    动态节点模式
    
    并行模式
    """
    args = argparse.ArgumentParser()
    args.add_argument('--config_path', required=False, help='path to case configs.json', default="smoke_case_config.json")
    args.add_argument('--image_name', required=False, help='image name with repository', default="")
    args.add_argument('--output_path', required=False, help='result output path', default="smoke_framework_output")
    args.add_argument('--config_root_dir', required=False, help='config root dir, within version_dir', default="")
    args.add_argument('--ver_conf_template', required=False, help='config template name, template name in all version_dir', default="")
    args.add_argument('--parallel', required=False, help='parallel number to run cases, default is 1', default=1, type=int)
    args.add_argument('--dynamic_nodes', action='store_true', help="Dynamically allocate free nodes")
    args.add_argument('--machine_pool_conf', required=False, help='machine_pool config path, when dynamic_nodes mode valid', default="./npu_manager/machine_pool.json")
    args.add_argument('--machine_pool_id', required=False, help='machine_pool id', default="blue_rdv_a3")
    args.add_argument('--version_mode', action='store_true', help="get config by version")
    args = args.parse_args()

    if not args.image_name:
        print('image: {args.image_name} and tag: {args.image_tag} are required')
        exit(-1)

    _config = Config(args.config_path, args.config_root_dir, args.ver_conf_template, args.version_mode)
    _node_manager = None
    if args.dynamic_nodes:
        # 动态节点模式，节点信息动态修改
        if not args.machine_pool_conf:
            print("dynamic nodes config are required")
            exit(-1)
        _node_manager = NodeManagement(args.machine_pool_conf, args.machine_pool_id)
    framework = SmokeFramework(_config, args.image_name, args.output_path, args.parallel,
                               args.dynamic_nodes, _node_manager)
    ret_code = framework.run()
    exit(ret_code)